import re
import json
import ast
from typing import Any, Dict, Optional

# ─────────────────────────────────────────────────────────────────────────────
# Helpers
# ─────────────────────────────────────────────────────────────────────────────

def _strip_code_fences(s: str) -> str:
    # ```json ... ``` 또는 ``` ... ``` 제거
    return re.sub(r"```[\w]*\n?|```", "", s, flags=re.MULTILINE).strip()

def _find_first_json_object(s: str) -> str | None:
    """
    스택으로 첫 번째 완결된 { ... } 블록 찾기.
    코드펜스 제거 후 사용 권장.
    """
    start = None
    depth = 0
    for i, ch in enumerate(s):
        if ch == "{":
            if depth == 0:
                start = i
            depth += 1
        elif ch == "}":
            if depth > 0:
                depth -= 1
                if depth == 0 and start is not None:
                    return s[start : i + 1]
    return None

def _find_last_json_object(s: str) -> Optional[str]:
    """
    문자열에서 마지막으로 완결된 최상위 JSON 오브젝트({ ... })를 찾아 반환.
    - 따옴표 내부의 { } 는 무시
    - 이스케이프 처리(\") 고려
    - 중첩 {} 지원
    """
    in_string = False
    escape = False
    depth = 0
    last_start = None
    last_end = None
    current_start = None

    for i, ch in enumerate(s):
        if in_string:
            if escape:
                escape = False
            elif ch == '\\':
                escape = True
            elif ch == '"':
                in_string = False
            continue
        else:
            if ch == '"':
                in_string = True
                continue
            if ch == '{':
                if depth == 0:
                    current_start = i
                depth += 1
                continue
            if ch == '}':
                if depth > 0:
                    depth -= 1
                    if depth == 0 and current_start is not None:
                        last_start = current_start
                        last_end = i + 1
                        current_start = None
                continue

    if last_start is not None and last_end is not None:
        return s[last_start:last_end]
    return None

def _soft_json_cleanup(s: str) -> str:
    """
    가벼운 정리: 트레일링 콤마 제거, BOM/제어문자 제거 등
    """
    s = s.replace("\ufeff", "")
    # },] 앞의 트레일링 콤마 제거
    s = re.sub(r",(\s*[}\]])", r"\1", s)
    # 공백 라인 과다 축소
    s = re.sub(r"\n{3,}", "\n\n", s)
    return s.strip()

def _safe_eval_expr(expr: str, data: Dict[str, Any]) -> Any:
    """
    {{ ... }} 안의 표현식을 안전하게 평가.
    허용: j["key"]["sub"][0] 형태의 인덱싱만.
    """
    tree = ast.parse(expr, mode="eval")

    def _eval(node, ctx):
        if isinstance(node, ast.Expression):
            return _eval(node.body, ctx)
        if isinstance(node, ast.Name):
            if node.id != "j":
                raise ValueError("Only 'j' is allowed.")
            return ctx["j"]
        if isinstance(node, ast.Subscript):
            base = _eval(node.value, ctx)
            sl = node.slice
            # Python 3.10+ 호환: slice가 직접 Constant / Tuple일 수 있음
            if isinstance(sl, ast.Constant):
                key = sl.value
            elif isinstance(sl, ast.Tuple):
                key = tuple(_eval(elt, ctx) for elt in sl.elts)
            else:
                # Python 3.8/3.9에서 ast.Index 사용 가능성
                if hasattr(ast, "Index") and isinstance(sl, ast.Index):
                    inner = sl.value
                    if isinstance(inner, ast.Constant):
                        key = inner.value
                    else:
                        raise ValueError("Unsupported index type.")
                else:
                    raise ValueError("Unsupported slice type.")
            return base[key]
        if isinstance(node, ast.Constant):
            return node.value
        # 그 외는 금지 (함수 호출, 연산 등)
        raise ValueError("Disallowed expression elements.")

    return _eval(tree, {"j": data})

def _render_template(tpl: str, data: Dict[str, Any]) -> str:
    """
    {{ ... }} 플레이스홀더를 _safe_eval_expr로 치환.
    """
    def repl(m: re.Match) -> str:
        expr = m.group(1).strip()
        try:
            val = _safe_eval_expr(expr, data)
        except Exception as e:
            val = f"[ERR:{e}]"
        return "" if val is None else str(val)

    return re.sub(r"\{\{\s*(.*?)\s*\}\}", repl, tpl)

DEFAULT_TEMPLATE = (
    'Subject & Attire: {{ j["Main Subject"] }}. Wearing {{ j["Initial Attire"] }}.\n'
    'Setting: {{ j["Background"] }}.\n'
    'Shot Plan (~3s each): 1) {{ j["Actions"][0] }} 2) {{ j["Actions"][1] }} 3) {{ j["Actions"][2] }}\n'
    'Camera: {{ j["Camera Composition"] }}.\n'
    'Final Look: {{ j["Final State"] }}.\n'
    'Style: {{ j["Style"] }}.\n'
)

# ─────────────────────────────────────────────────────────────────────────────
# Node 1: 텍스트 출력에서 JSON 추출 (Text -> Text(JSON))
# ─────────────────────────────────────────────────────────────────────────────
class ExtractJSON:
    """
    입력: text (LLM 원본 출력)
    옵션:
      - strategy: 'last' 또는 'first' (기본 last)
      - fallback_to_raw: 파싱 실패 시 정리된 원문 반환 여부
    출력:
      - json (STRING): 미니파이된 JSON 문자열
    우측 프리뷰:
      - ui.text: pretty-printed JSON (indent=2), 실패 시 원문
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "text": ("STRING", {"multiline": True}),
                "strategy": (["last", "first"], {"default": "last"}),
                "fallback_to_raw": ("BOOLEAN", {"default": True}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("json",)
    FUNCTION = "extract"
    CATEGORY = "PromptTools/Qwen3"

    def extract(self, text: str, strategy: str = "last", fallback_to_raw: bool = True):
        s = _strip_code_fences(text or "")

        if strategy == "last":
            candidate = _find_last_json_object(s)
            if candidate is None:
                # 마지막이 없으면 전체 문자열에서라도 시도 (보수적 폴백)
                candidate = s
        else:  # "first"
            candidate = _find_first_json_object(s) or s

        cleaned = _soft_json_cleanup(candidate)

        try:
            obj = json.loads(cleaned)
            # 포트 출력: 미니파이 JSON (파이프 안정)
            out = json.dumps(obj, ensure_ascii=False, separators=(",", ": "))
            # 프리뷰: 예쁘게
            ui_str = json.dumps(obj, ensure_ascii=False, indent=2)
        except Exception:
            out = cleaned if fallback_to_raw else ""
            ui_str = out  # 프리뷰도 원문

        return (out, {"ui": {"text": (ui_str,)}})
# ─────────────────────────────────────────────────────────────────────────────
# Node 2: 템플릿 렌더링 (내부 text input 포함) (Text(JSON) -> Text)
# ─────────────────────────────────────────────────────────────────────────────
class TemplateRenderFromJSON:
    """
    입력:
      - json_text: 문자열(JSON)
      - template: 내부 텍스트 입력에 템플릿을 직접 작성
        (플레이스홀더: {{ j["Main Subject"] }}, {{ j["Actions"][0] }} 등)
    출력:
      - text: 완성된 최종 프롬프트 문자열
    """
    @classmethod
    def INPUT_TYPES(cls):
        return {
            "required": {
                "json_text": ("STRING", {"multiline": True}),
                "template": ("STRING", {"multiline": True, "default": DEFAULT_TEMPLATE}),
            }
        }

    RETURN_TYPES = ("STRING",)
    RETURN_NAMES = ("text",)
    FUNCTION = "render"
    CATEGORY = "PromptTools/json"

    def render(self, json_text: str, template: str):
        try:
            data = json.loads(json_text)
        except Exception as e:
            # JSON 파싱 실패 시, 빈 dict로 처리하고 오류 표시
            data = {}
            template = f"[JSON Parse Error: {e}]\n{template}"
        text = _render_template(template, data)

        ui = {"ui": {"text": [text]}}
        return (text, ui)

# ─────────────────────────────────────────────────────────────────────────────
# ComfyUI registry
# ─────────────────────────────────────────────────────────────────────────────
NODE_CLASS_MAPPINGS = {
    "ExtractJSON": ExtractJSON,
    "TemplateRenderFromJSON": TemplateRenderFromJSON,
}

NODE_DISPLAY_NAME_MAPPINGS = {
    "ExtractJSON": "Extract JSON",
    "TemplateRenderFromJSON": "Prompt: Render Template from JSON",
}
